; 32-bit cooperative context switch helper
; Exports: ctx_switch(oldEspPtr, newEsp)
; Saves current ESP into *oldEspPtr, loads new ESP, and returns.
; When both contexts call a function like yield() which invokes ctx_switch,
; ret will resume execution in the other context right after its yield() call.

BITS 32
global ctx_switch
global ctx_make_taskA
global ctx_make_taskB
global ctx_make_taskElf
global ctx_make_taskWorker
extern taskA_entry
extern taskB_entry
extern elfTask_entry
extern worker_entry

section .text

; void ctx_switch(uint32_t* oldEspPtr, uint32_t newEsp)
ctx_switch:
    push ebp
    mov  ebp, esp
    ; [ebp+8]  = oldEspPtr
    ; [ebp+12] = newEsp
    mov  eax, [ebp+8]    ; eax = oldEspPtr
    mov  edx, esp        ; edx = current ESP (pointing to saved EBP)
    mov  [eax], edx      ; *oldEspPtr = current ESP
    mov  edx, [ebp+12]   ; edx = newEsp
    mov  esp, edx        ; switch to new stack
    pop  ebp             ; restore callee frame from new stack
    ret                  ; return into the new context's caller

; Prepare initial stack for task A: place return address = taskA_entry
; uint32_t ctx_make_taskA(uint32_t top)
ctx_make_taskA:
    push ebp
    mov  ebp, esp
    mov  eax, [ebp+8]    ; eax = top
    and  eax, 0xFFFFFFF0 ; 16-byte align
    sub  eax, 4          ; reserve space for return address
    mov  dword [eax], taskA_entry
    pop  ebp
    ret

; Prepare initial stack for task B: place return address = taskB_entry
; uint32_t ctx_make_taskB(uint32_t top)
ctx_make_taskB:
    push ebp
    mov  ebp, esp
    mov  eax, [ebp+8]    ; eax = top
    and  eax, 0xFFFFFFF0 ; 16-byte align
    sub  eax, 4          ; reserve space for return address
    mov  dword [eax], taskB_entry
    pop  ebp
    ret

; Prepare initial stack for ELF task: place return address = elfTask_entry
; uint32_t ctx_make_taskElf(uint32_t top)
ctx_make_taskElf:
    push ebp
    mov  ebp, esp
    mov  eax, [ebp+8]    ; eax = top
    and  eax, 0xFFFFFFF0 ; 16-byte align
    sub  eax, 4          ; reserve space for return address
    mov  dword [eax], elfTask_entry
    pop  ebp
    ret

; Prepare initial stack for generic worker task: return address = worker_entry
; uint32_t ctx_make_taskWorker(uint32_t top)
ctx_make_taskWorker:
    push ebp
    mov  ebp, esp
    mov  eax, [ebp+8]    ; eax = top
    and  eax, 0xFFFFFFF0 ; 16-byte align
    sub  eax, 4          ; reserve space for return address
    mov  dword [eax], worker_entry
    pop  ebp
    ret